import asyncio
import io
import time
import traceback
from datetime import datetime

import cv2
import numpy as np
from PIL import Image
import io
import wda
from fastapi import WebSocket, APIRouter, Depends
from typing import List

from loguru import logger

from apps.system.cruds.agent import agent_dao
from apps.system.cruds.device import device_dao
from apps.system.schema.agent import UpdateAgent
from database.db import async_db_session

# import scrcpy
# from adbutils import adb

router = APIRouter(prefix="/server/websockets/agent", tags=["agent_websocket"])


class ConnectionManager:
    def __init__(self):
        self.active_connections: List[WebSocket] = []
        # self.driver: wda.USBClient = None
        # self.driver: scrcpy.Client = None
        self.queue = None

    async def connect(self, websocket: WebSocket, secret_key: str):
        # BytesTool.agentId = jsonObject.getInteger("id");
        # BytesTool.highTemp = jsonObject.getInteger("highTemp");
        # BytesTool.highTempTime = jsonObject.getInteger("highTempTime");
        # BytesTool.remoteTimeout = jsonObject.getInteger("remoteTimeout");
        # info = {'msg': 'agentInfo', 'agentId': 1, 'port': 7777, 'systemType': 'Mac OS X', 'host': '127.0.0.1', 'version': 'v1.0.0', 'hasHub': 0}
        await websocket.accept()
        async with async_db_session() as db:
            agent = await agent_dao.get_agent_by_secret_key(db, secret_key)
            if agent:
                await websocket.send_json({
                    'msg': 'auth', 'result': 'pass',
                    'id': agent.id, 'highTemp': agent.high_temp, 'highTempTime': agent.high_temp_time, 'remoteTimeout': 900
                })
                self.active_connections.append(websocket)

    def disconnect(self, websocket: WebSocket):
        self.active_connections.remove(websocket)

    async def broadcast(self, message: str):
        for connection in self.active_connections:
            await connection.send_text(message)


manager = ConnectionManager()


@router.websocket("/{secret_key}")
async def websocket_endpoint(websocket: WebSocket, secret_key):
    await manager.connect(websocket, secret_key)
    # print(secret_key)
    try:
        while True:
            # await websocket.send_json({'msg': 'pong'})
            # await asyncio.sleep(1)
            msg = await websocket.receive_json()
            match msg['msg']:
                case "agentInfo":
                    async with async_db_session() as db:
                        agent = {
                            'host': msg['host'],
                            'port': msg['port'],
                            'system_type': msg['systemType'],
                            'version': msg['version'],
                            'has_hub': msg['hasHub']
                        }
                        await agent_dao.update_agent(db, msg['agentId'], agent)
                case "ping":
                    await websocket.send_json({"msg": "pong"})
                case "deviceDetail":
                    if msg['udId'] != '' and 'Error' not in msg['udId']:
                        keys = {
                            'agentId': 'agent_id',
                            'udId': 'ud_id',
                            'isHm': 'is_hm'
                        }
                        msg.pop('msg')
                        for key, value in keys.items():
                            if key in msg:
                                msg.update({value: msg.pop(key)})
                        logger.info(msg)
                        async with async_db_session() as db:
                            await device_dao.create_or_update(db, msg)


    # {'msg': 'deviceDetail', 'agentId': 1, 'size': '1080x2376', 'udId': 'XWN0221727000549', 'platform': 1, 'isHm': 0, 'status': 'OFFLINE'}
    # {'msg': 'deviceDetail', 'agentId': 1, 'size': '1080x2376', 'udId': 'XWN0221727000549', 'platform': 1, 'isHm': 0}
    # {'msg': 'deviceDetail', 'agentId': 1, 'size': '1080x2376', 'name': 'OCE-AN10', 'cpu': 'arm64-v8a', 'model': 'OCE-AN10', 'udId': 'XWN0221727000549', 'version': '4.2.0', 'platform': 1, 'isHm': 1, 'status': 'ONLINE', 'manufacturer': 'HUAWEI'}
    # {'msg': 'battery', 'agentId': 1, 'detail': [{'vol': 4203, 'level': 94, 'udId': 'XWN0221727000549', 'tem': 170}]}
    except Exception as e:
        traceback.print_exc()
        print("关闭连接", e)
        # manager.disconnect(websocket)
